#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/4/15 14:42
# @Author  : hqu-jcl
# @Site    : 
# @File    : HistogramEqualization.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
import copy

"""
    直方图均衡
"""


class HistogramEqualization:
    def __init__(self, _img):
        self.image = _img

    def getImage(self):
        """
        :return:  图片矩阵
        """
        return self.image

    def display(self):
        """
         显示图片
        :param self:
        :return:
        """
        temp = np.array(self.image, dtype=np.uint8)
        cv2.imshow('image', temp)
        cv2.waitKey(0)

    def getHistogram(self):
        """
        获取频率直方图
        :param self:
        :return:   频率直方图
        """
        hist = np.zeros(256, np.int)
        for k in range(len(self.image)):
            for j in range(len(self.image[0])):
                hist[int(self.image[k, j])] += 1
        return hist

    def plotHistogram(self):
        """
         绘制直方图
        :param self:
        :return:
        """
        n = len(self.image[0]) * len(self.image)  # 总像素点个数
        hist = self.getHistogram()
        group = range(0, 256)
        plt.bar(group, hist)
        plt.show()

    def equalization(self):
        """
        直方图均衡
        :param self:
        :return: 均衡后的图像矩阵
        """
        _hist = self.getHistogram()
        # 总像素点
        n = len(self.image[0]) * len(self.image)
        l = 256  # 灰度级个数
        hist_prob = _hist / n
        S = []  # 累积分布概率
        for i in range(len(hist_prob)):
            S.append(sum(hist_prob[:i + 1]))
        print(S)
        new_pixel = []
        for i in range(len(S)):
            new_pixel.append(round(S[i] * (l - 1)))
        # 四舍五入 合并

        #  映射像素点
        _img = np.zeros([len(self.image), len(self.image[0])])
        for i in range(len(self.image)):
            for j in range(len(self.image[0])):
                _img[i, j] = new_pixel[int(self.image[i, j])]
        return _img


if __name__ == '__main__':
    img = cv2.imread("../resource/f4.png", 0)
    img = np.array(img, np.uint8)
    HE = HistogramEqualization(img)
    HE.plotHistogram()
    HE.display()
    img1 = HE.equalization()

    # img1 = HistogramEqualization(img1)
    # img1.plotHistogram()  # 绘制直方图
    # img1.display()  # 结果

    # 测试多次均衡化
    temp = np.array(img, dtype=np.uint8)
    for i in range(5):
        _temp1 = HistogramEqualization(temp)
        _temp = _temp1.equalization()
        temp = HistogramEqualization(_temp)
        temp.plotHistogram()
        temp.display()
        temp = temp.getImage()

    # hist=cv2.calcHist([equ],[0],None,[256],[0,255])
